#!/usr/bin/env python
"""
Copyright Ian Ross Williams, 2012

    This file is part of ArcSecond.

    ArcSecond is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    ArcSecond is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with ArcSecond.  If not, see <http://www.gnu.org/licenses/>.
"""
import numpy as np
from numpy.fft import rfft2,irfft2
import Image,ImageOps
from scipy.optimize import fmin_powell
import scipy.interpolate as interpolate
from fipy import *
from arcsecond import filters
import pywt
import os
import time

debugPath=None
if os.path.exists("/home/ian/workspace/gtkapture/"):
    debugPath = "/home/ian/workspace/gtkapture/"
elif os.path.exists("/home/ian/Code/gtkapture/"):
    debugPath = "/home/ian/Code/gtkapture/"
elif os.path.exists("/home/wi10008/Code/gtkapture/gtkapture/"):
    debugPath = "/home/wi10008/Code/gtkapture/gtkapture/"

class MFBC:
    ( INITIAL_FIRSTFRAME,
      INITIAL_MEAN,
      INITIAL_MEDIAN,
      INITIAL_KALMAN) = range(4)
           
    #Shapelet velocity Multi-Frame Blind Deconvolution
    def MFBC(self,stack,shape,
             stackExtra=None,callback=None,roi=None,
             tolerance=1./255.,underSample=10,distance=10.,
             globalIterations=2,fittingIterations=3.,
             methodInitialGuess=INITIAL_FIRSTFRAME,
             ignoreImagesWithStdDevBelow=1.,noOfFrames=1):

        if debugPath!=None:
            os.system("rm " + debugPath + "src/test/*")
        fittingIterations=int(fittingIterations)
        globalIterations=int(globalIterations)
        noOfFrames=int(noOfFrames)
        underSample=int(underSample)
        
        allFrames = stack.shape[0]
        mask = np.arange(stack.shape[0]) < noOfFrames
        sstack = stack[mask]
        if stackExtra!=None:
            for channelStack in stackExtra: 
                channelStack = channelStack[mask]
            
        tolerance = max(tolerance,1./(np.max(stack) - np.min(stack)))
        
        if roi==None: roi=(0,0,shape[1],shape[0])
        roi = (int(roi[0]),int(roi[1]),int(roi[2]),int(roi[3]))
        x0,y0=roi[1],roi[0]
        x1,y1=roi[3],roi[2]
        sampleshape=np.array((x1-x0,y1-y0))
        print roi, sampleshape

        #Create the velocity field for the image sample
        velocityshape = sampleshape/underSample
        velocityshape[velocityshape<5]=5
        velocitymx, velocitymy = (x1+x0)/2,(y1+y0)/2
        velocityx0 = velocitymx - velocityshape[0]/2
        velocityy0 = velocitymy - velocityshape[1]/2
        velocityx1 = velocityx0 + velocityshape[0]
        velocityy1 = velocityy0 + velocityshape[1]
        vstack = np.zeros((noOfFrames, 2,velocityshape[0],velocityshape[1]))
        
        #Define the velocity bounds
        nb=2*np.prod(velocityshape)
        bounds = np.array( ([-1.,1.])*nb )
        bounds = np.reshape(bounds,(nb,2))
        
        #Form an initial guess at the answer
        xn = sstack[0]
        xn = xn.reshape(shape)
        imX = Image.fromarray(xn.astype(np.int8),"L")
        
        if methodInitialGuess==self.INITIAL_MEAN:
            alive=callback(imX, "Forming initial guess using Mean stack" )
            if not alive: return None
            xn = np.mean(stack,axis=0)
            xn = xn.reshape(shape)
        if methodInitialGuess==self.INITIAL_MEDIAN:
            alive=callback(imX, "Forming initial guess using Median stack" )
            if not alive: return None
            xn = np.median(stack,axis=0)
            xn = xn.reshape(shape)
        elif methodInitialGuess==self.INITIAL_KALMAN:
            alive=callback(imX, "Forming initial guess using Kalman filter" )
            if not alive: return None

            xhat = np.median(stack, axis=0)   #A posteriori estimate of x
            P         = np.ones(xhat.shape)   #A posteriori error estimate
            xhatMinus = np.zeros(xhat.shape)  #A priori estimate of x
            Pminus    = np.zeros(xhat.shape)  #A priori error estimate
            K         = np.zeros(xhat.shape)  #Gain or blending factor
            error=np.std(stack,axis=0)
            Q = np.mean(error) #Process noise
            R = np.var(error)  #Estimate of measurement error variance
            for frame in stack:
                xhatMinus[:] = xhat
                Pminus[:] = P + Q
                K[:] = Pminus / (Pminus + R)
                xhat[:] = xhatMinus + K*(frame-xhatMinus)
                P[:] = (1. - K) * Pminus
            xn = xhat
            xn = xn.reshape(shape)

        #Store the initial guess for reference
        xn0 = np.array(xn)
        imX = Image.fromarray(xn0.astype(np.int8),"L")
        if debugPath!=None:
            imX.save("" + debugPath + "src/test/x0.png")

        #Estimate the initial noise as y = p*x + n => n = y - p*x
        print "Noise estimate:",
        noiseMean = np.zeros(shape)
        noiseMean = np.fft.rfft2(noiseMean,s=shape)
        noiseSigma=0
        for i in xrange(noOfFrames):
            yt = np.reshape(sstack[i],shape)
            ytn = yt / np.mean(yt) * np.mean(xn)
            noise = ytn - xn
            noiseMean = noiseMean + np.abs(np.fft.rfft2(noise,s=shape))
            noiseSigma = noiseSigma + np.sqrt(np.mean(noise**2))
        noiseSpectra = noiseMean / noOfFrames
        noiseSpectra = noiseSpectra / np.max(noiseSpectra)
        noiseSigma = noiseSigma / noOfFrames
        NSR = noiseSigma**2/np.mean(xn**2) 
        print "Noise to Signal Ratio=%.1f%%" % (NSR*100.)
        imX = Image.fromarray(xn.astype(np.int8),"L")
        alive=callback(imX, "Noise to Signal Ratio=%.1f%%" % (NSR*100.) )
        if not alive: return None      
                            
        #Save the input tile
        for i,ytflat in enumerate(sstack):
            yt = np.reshape(ytflat,shape)
            sampleyt = yt[x0:x1,y0:y1]
            imSY = Image.fromarray(sampleyt.astype(np.int8),"L")
            if debugPath!=None:
                imSY.save("" + debugPath + "src/test/sy%03d.png" % i)
                              
        entropyxn0 = -np.mean(xn0[xn0>0.]*np.log(xn0[xn0>0.]))
        sharpnessxn0 = np.mean(np.array(np.gradient(xn0))**2)
        
        #Fitting residual
        def residual(velocity):
            vres = np.reshape(velocity,(2,velocityshape[0],velocityshape[1]))
            #Convect the current frame
            convected = self.convect(sampleyt,vres[0],vres[1],distance )
            convected = np.clip(convected,0.,255.)
                        
            #Score the convection by correlation with the best guess image
            beta = 3.    #The third norm or greater avoids unstable convection solutions
            resid = np.mean( np.abs((samplex[distance:-distance,distance:-distance] - convected[distance:-distance,distance:-distance]))**beta)* (1./beta)
            
            if callback!=None:
                #xnH = np.array(convected)
                #imX = Image.fromarray(xnH.astype(np.int8),"L")
                msg = "Resid=%f" % (resid)
                #alive=callback(imX, msg )
                print msg
                if not alive: return 0
            
            if 1:
                #Bias the velocity choice to enhance a sharpness measure
                #sharpness = np.mean(np.array(np.gradient(convected))**2)
                #resid = resid - sharpness/sharpnessxn0 
                entropy = -np.mean(convected[convected>0.]*np.log(convected[convected>0.]))
                resid = resid - entropy/entropyxn0

            return resid
        
        for EMiter in xrange(globalIterations):
            #Retrieve the current solution tile
            samplex  = xn[x0:x1,y0:y1]
            samplex0 = xn0[x0:x1,y0:y1]
            cstack = np.zeros((noOfFrames,sampleshape[0],sampleshape[1]))
            for i in xrange(noOfFrames):
                #Extract the current frame from the sample stack
                yt = np.array(np.reshape(sstack[i],shape))[x0:x1,y0:y1]
                                    
                #Fit the velocity for the current frame
                vi = vstack[i].flatten()
                #vi, val,d = fmin_l_bfgs_b(residual,vi,approx_grad=True,bounds=bounds,epsilon=0.01,maxfun=fittingIterations)
                vi = fmin_powell(residual,vi,maxiter=fittingIterations,ftol=0.01)
                if not alive: return None
                vi = np.reshape(vi, (2,velocityshape[0],velocityshape[1]))
                vstack[i] = vi

                #Convect the current frame
                vi = vstack[i]
                convected = self.convect(yt,vi[0],vi[1],distance )
                print yt.shape, convected.shape, cstack[i].shape
                cstack[i] = convected
                cstack[i] = np.clip(cstack[i],0.,255.)
                    
                #Display progress
                if callback!=None:
                    xnH = np.array(cstack[i])
                    imX = Image.fromarray(xnH.astype(np.int8),"L")
                    msg = "Iteration %d. velocity for frame %d of %d" % (EMiter+1, i+1, noOfFrames)
                    alive=callback(imX, msg )
                    print msg
                    if not alive: return None

            #Form the next estimate
            xnNew = np.median(cstack,axis=0)
            print np.prod(sampleshape),xn.shape
            xnNew = np.reshape(xnNew,sampleshape)
            deltax = xnNew-xn[x0:x1,y0:y1]
            xn[x0+distance:x1-distance,y0+distance:y1-distance] = xnNew[distance:-distance,distance:-distance]
            
            #Estimate the noise level
            print "Noise estimate:",
            noiseMean = np.zeros(shape)
            noiseMean = np.fft.rfft2(noiseMean,s=shape)
            noiseSigma=0
            for i in xrange(noOfFrames):
                yt = np.reshape(sstack[i],shape)
                ytn = yt / np.mean(yt) * np.mean(xn)
                noise = ytn - xn
                noiseMean = noiseMean + np.abs(np.fft.rfft2(noise,s=shape))
                noiseSigma = noiseSigma + np.sqrt(np.mean(noise**2))
            noiseSpectra = noiseMean / noOfFrames
            noiseSpectra = noiseSpectra / np.max(noiseSpectra)
            noiseSigma = noiseSigma / noOfFrames
            
            NSR = noiseSigma**2/np.mean(xn**2) 
            print "Noise to Signal Ratio=%.1f%%" % (NSR*100.)

            globalResidual = np.sqrt(np.max(deltax**2))/np.max(xn)
            print "Global residual=", globalResidual
            if globalResidual<tolerance: break
            
            #Store the result image
            xClipped = np.clip(xn, 0,255)
            imX = Image.fromarray(xClipped.astype(np.int8),"L")
            if debugPath!=None:
                imX.save("" + debugPath + "src/test/x.png")
                        
            if callback!=None:
                alive=callback(imX, "Iteration %d.  Max change=%.1f%%.  NSR=%.2f%%" % (EMiter+1,globalResidual*100.,NSR*100.) )
                if not alive: return None
            
        if callback!=None:
            alive=callback(imX, "MFBD Complete. Result saved to x.png" )
            if not alive: return None

        return xn

    def convect(self, sample, velocityx,velocityy,distance=10.,CFL=1.0):
        #Interpolate the velocity field to the same size as the sample
        nx,ny = sample.shape[0],sample.shape[1]
        vnx,vny = velocityx.shape[0],velocityx.shape[1]
        nxx,nyy = np.arange(0.5,nx+0.5),np.arange(0.5, ny+0.5)
        rx,ry = float(nx)/float(vnx),float(ny)/float(vny)
        vxx,vyy = np.arange(rx/2.,nx+rx/2.,rx),np.arange(ry/2.,ny+ry/2.,ry)
        ufx = interpolate.RectBivariateSpline(vxx,vyy,velocityx)
        ufy = interpolate.RectBivariateSpline(vxx,vyy,velocityy)
        uxFit = np.array(ufx(nxx,nyy)/nx)
        uyFit = np.array(ufy(nxx,nyy)/ny)
        uxFit = np.array(uxFit.flatten(),dtype=np.float32)
        uyFit = np.array(uyFit.flatten(),dtype=np.float32)

        #Choose the time-step assuming velocities are bounded by -1 and +1.
        umax = np.sqrt(np.max(uxFit**2 + uyFit**2))
        dt = CFL/max(1e-6,umax)
        steps=max(1,int(distance/dt))
        
        #Create the mesh and variables
        mesh = Grid2D(nx=nx, ny=ny)
        phi = CellVariable(name  = "image data",
                           mesh  = mesh,
                           value = 0.,
                           hasOld= 1)
        ux = CellVariable(mesh=mesh, name='X velocity',value=0.)
        uy = CellVariable(mesh=mesh, name='Y velocity',value=0.)
        u = FaceVariable(name= "velocity",mesh=mesh, rank=1,value=1.)

        #Initial condition
        image = np.array(sample.flatten(),dtype=np.float32)
        phi.setValue( image )

        ux.setValue( uxFit )
        uy.setValue( uyFit )
        u[0] = ux.getArithmeticFaceValue()
        u[1] = uy.getArithmeticFaceValue()

        #Set the boundary conditions 
        faces = (mesh.getFacesTop() | mesh.getFacesRight() | mesh.getFacesBottom() | mesh.getFacesLeft())
        BCs = FixedValue(faces=faces, value=phi)

        #Define the equation to solve
        #convection = CentralDifferenceConvectionTerm # fails
        convection = ExponentialConvectionTerm # (1)
        #convection = HybridConvectionTerm # (1)
        #convection = PowerLawConvectionTerm # (1)
        #convection = UpwindConvectionTerm # (2)
        #convection = ExplicitUpwindConvectionTerm # (3)
        #convection = VanLeerConvectionTerm # (3)

        eq = TransientTerm(coeff=1.) == convection(coeff=u)

        #Convect
        for step in range(steps):
            phi.updateOld()
            eq.solve(phi, dt=dt)
            
        #Extract the solution
        nim = np.array(phi,dtype=np.float32)
        nim = np.reshape(nim,sample.shape)
        return nim
